"""
MCP插件加载器
支持动态加载和管理MCP插件
"""

import os
import importlib.util
import json
import logging
import requests
import zipfile
import shutil
import hashlib
import hmac
import asyncio
import functools
from typing import Dict, Any, List, Optional, Callable, Set
import aiohttp
from src.config.settings import settings
from src.services.mcp_service import mcp_service_manager


class MCPPluginError(Exception):
    """MCP插件相关异常"""
    def __init__(self, message: str, details: Optional[Dict[str, Any]] = None):
        super().__init__(message)
        self.message = message
        self.details = details or {}


class MCPPluginLoadError(MCPPluginError):
    """MCP插件加载异常"""
    pass


class MCPPluginDependencyError(MCPPluginError):
    """MCP插件依赖异常"""
    pass


class MCPPluginInstallError(MCPPluginError):
    """MCP插件安装异常"""
    pass


class MCPPluginSecurityError(MCPPluginError):
    """MCP插件安全异常"""
    pass


class MCPPluginStateError(MCPPluginError):
    """MCP插件状态异常"""
    pass


class MCPPluginLoader:
    """MCP插件加载器"""
    
    # 插件状态常量
    PLUGIN_STATE_LOADED = "loaded"
    PLUGIN_STATE_INITIALIZED = "initialized"
    PLUGIN_STATE_PAUSED = "paused"
    PLUGIN_STATE_RUNNING = "running"
    PLUGIN_STATE_STOPPED = "stopped"
    PLUGIN_STATE_ERROR = "error"
    
    def __init__(self, plugin_dir: str = "mcps"):
        self.plugin_dir = plugin_dir
        self.loaded_plugins = {}
        self.plugin_states = {}  # 插件状态管理
        self.plugin_lifecycle_hooks = {}  # 插件生命周期钩子
        self.plugin_dependencies = {}  # 插件依赖关系
        self.resolved_dependencies = {}  # 已解析的依赖
        self.plugins_config_file = os.path.join(plugin_dir, "plugins_config.json")
        self.plugin_registry_url = "https://raw.githubusercontent.com/langchain-graph/mcp-registry/main"  # 示例插件注册表URL
        self.plugin_secret_key = os.getenv("MCP_PLUGIN_SECRET_KEY", "")  # 插件签名密钥
        self.logger = logging.getLogger(__name__)
        self._plugin_cache: Dict[str, Any] = {}  # 初始化插件缓存
        self.dependency_graph: Dict[str, List[str]] = {}  # 依赖图

    def discover_plugins(self) -> List[str]:
        """
        发现可用的插件
        
        Returns:
            插件名称列表
        """
        try:
            plugins = []
            if os.path.exists(self.plugin_dir):
                for item in os.listdir(self.plugin_dir):
                    item_path = os.path.join(self.plugin_dir, item)
                    if os.path.isdir(item_path) and not item.startswith('__'):
                        plugins.append(item)
            return plugins
        except Exception as e:
            error_details = {
                "plugin_dir": self.plugin_dir,
                "error_type": type(e).__name__,
                "error_message": str(e)
            }
            self.logger.error(f"发现插件失败: {error_details}")
            raise MCPPluginLoadError(f"发现插件失败: {str(e)}", error_details) from e

    def load_plugin_from_directory(self, plugin_name: str) -> Dict[str, Any]:
        """
        从目录加载插件配置
        
        Args:
            plugin_name: 插件名称
            
        Returns:
            插件配置
        """
        plugin_path = os.path.join(self.plugin_dir, plugin_name)
        if not os.path.exists(plugin_path):
            raise FileNotFoundError(f"插件目录 {plugin_path} 不存在")
        
        # 查找配置文件
        config_file = os.path.join(plugin_path, "config.json")
        if not os.path.exists(config_file):
            # 尝试从Dockerfile或其他配置文件中提取信息
            return self._extract_config_from_files(plugin_path, plugin_name)
        
        with open(config_file, 'r', encoding='utf-8') as f:
            config = json.load(f)
            
        # 确保配置中包含版本信息
        if "version" not in config:
            config["version"] = "1.0.0"  # 默认版本
            
        # 处理依赖关系
        if "dependencies" in config:
            self.plugin_dependencies[plugin_name] = config["dependencies"]
            self.dependency_graph[plugin_name] = config["dependencies"]
            
        return config
    
    def _extract_config_from_files(self, plugin_path: str, plugin_name: str) -> Dict[str, Any]:
        """
        从文件中提取配置
        
        Args:
            plugin_path: 插件路径
            plugin_name: 插件名称
            
        Returns:
            提取的配置
        """
        config = {
            "server_name": plugin_name,
            "enabled": True,
            "version": "1.0.0",  # 默认版本
            "dependencies": []   # 默认无依赖
        }
        
        # 尝试从README文件中提取版本信息
        readme_files = ["README.md", "README.txt", "README"]
        for readme in readme_files:
            readme_path = os.path.join(plugin_path, readme)
            if os.path.exists(readme_path):
                with open(readme_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                    # 简单的版本提取（查找类似 "version x.x.x" 的模式）
                    import re
                    version_match = re.search(r'[vV]ersion\s*:?\s*([0-9]+\.?[0-9]*\.?[0-9]*)', content, re.IGNORECASE)
                    if version_match:
                        config["version"] = version_match.group(1)
                        break
        
        # 尝试从compose.yml提取端口信息
        compose_file = os.path.join(plugin_path, "compose.yml")
        if os.path.exists(compose_file):
            # 简单解析compose文件以获取端口信息
            with open(compose_file, 'r', encoding='utf-8') as f:
                content = f.read()
                if "ports:" in content:
                    # 简化的端口提取
                    lines = content.split('\n')
                    for line in lines:
                        if "- " in line and ":" in line:
                            port_mapping = line.strip().split('- ')[1]
                            host_port = port_mapping.split(':')[0]
                            try:
                                config["url"] = f"http://localhost:{host_port}"
                                break
                            except:
                                pass
        
        # 如果没有从compose.yml获取到URL，使用默认值
        if "url" not in config:
            config["url"] = f"http://localhost:8000/{plugin_name}"
            
        return config

    def check_version_compatibility(self, plugin_name: str, required_version: str, actual_version: str) -> bool:
        """
        检查插件版本兼容性
        
        Args:
            plugin_name: 插件名称
            required_version: 所需版本
            actual_version: 实际版本
            
        Returns:
            是否兼容
        """
        try:
            from packaging import version
            
            # 简单的版本比较
            required = version.parse(required_version)
            actual = version.parse(actual_version)
            
            # 检查是否满足最低版本要求
            return actual >= required
        except Exception as e:
            self.logger.warning(f"版本兼容性检查失败: {e}")
            # 如果无法解析版本，假设兼容
            return True

    def update_plugin(self, plugin_name: str, target_version: str = "latest") -> bool:
        """
        更新插件到指定版本
        
        Args:
            plugin_name: 插件名称
            target_version: 目标版本
            
        Returns:
            是否更新成功
        """
        try:
            # 检查插件是否存在
            if plugin_name not in self.loaded_plugins:
                self.logger.error(f"插件 {plugin_name} 未注册")
                return False
            
            # 卸载当前版本
            self.unregister_plugin(plugin_name)
            
            # 从插件市场下载新版本（模拟）
            if target_version == "latest":
                self.logger.info(f"正在更新插件 {plugin_name} 到最新版本")
            else:
                self.logger.info(f"正在更新插件 {plugin_name} 到版本 {target_version}")
            
            # 重新注册插件
            success = self.register_plugin(plugin_name)
            
            if success:
                self.logger.info(f"插件 {plugin_name} 更新成功")
            else:
                self.logger.error(f"插件 {plugin_name} 更新失败")
                
            return success
        except Exception as e:
            self.logger.error(f"更新插件 {plugin_name} 失败: {e}")
            return False

    def _topological_sort(self, dependencies: Dict[str, List[str]]) -> List[str]:
        """
        对插件依赖进行拓扑排序，确保依赖项先于被依赖项加载
        
        Args:
            dependencies: 插件依赖关系图，键为插件名，值为依赖的插件列表
            
        Returns:
            按依赖顺序排序的插件列表
            
        Raises:
            MCPPluginDependencyError: 当存在循环依赖时抛出
        """
        # 构建完整的依赖图，确保所有被依赖的节点都包含在内
        complete_graph = {}
        # 首先添加所有直接定义的节点
        for node in dependencies:
            complete_graph[node] = list(dependencies[node])
        
        # 然后添加所有被依赖但未在顶层定义的节点
        for node in dependencies:
            for dep in dependencies[node]:
                if dep not in complete_graph:
                    complete_graph[dep] = []
        
        # 构建入度表（被依赖次数）
        in_degree = {}
        # 初始化所有节点的入度为0
        for node in complete_graph:
            in_degree[node] = 0
        
        # 计算每个节点的入度（被多少个节点依赖）
        for node in complete_graph:
            for dep in complete_graph[node]:  # node依赖dep
                in_degree[dep] += 1  # dep被node依赖，所以dep的入度加1

        # 找到所有入度为0的节点（不依赖任何其他节点的节点）
        queue = [node for node in in_degree if in_degree[node] == 0]
        result = []

        while queue:
            node = queue.pop(0)
            result.append(node)

            # 减少依赖该节点的所有节点的入度
            for dependent in [n for n in complete_graph if node in complete_graph[n]]:
                in_degree[dependent] -= 1
                if in_degree[dependent] == 0:
                    queue.append(dependent)

        # 检查是否有循环依赖
        if len(result) != len(complete_graph):
            raise MCPPluginDependencyError("检测到插件间存在循环依赖")

        return result

    def resolve_dependencies(self, plugin_name: str, plugin_config: Dict[str, Any]) -> List[str]:
        """
        解析插件依赖关系
        
        Args:
            plugin_name: 插件名称
            plugin_config: 插件配置
            
        Returns:
            按加载顺序排列的依赖列表
        """
        # 从插件配置中提取依赖信息
        dependencies = plugin_config.get("dependencies", [])
        
        # 构建依赖图
        dependency_graph = {plugin_name: dependencies}
        
        # 添加依赖项到依赖图
        for dep in dependencies:
            if dep in self.loaded_plugins:
                dep_config = self.loaded_plugins[dep]
                dependency_graph[dep] = dep_config.get("dependencies", [])
            else:
                # 尝试从插件目录加载依赖配置
                try:
                    dep_config = self.load_plugin_from_directory(dep)
                    dependency_graph[dep] = dep_config.get("dependencies", [])
                except Exception as e:
                    raise MCPPluginDependencyError(f"无法加载依赖插件 {dep}: {str(e)}")

        # 使用拓扑排序确定加载顺序
        try:
            sorted_deps = self._topological_sort(dependency_graph)
            # 移除当前插件本身（它应该在最后加载）
            sorted_deps.remove(plugin_name)
            return sorted_deps
        except MCPPluginDependencyError:
            raise
        except Exception as e:
            raise MCPPluginDependencyError(f"解析依赖关系时出错: {str(e)}")

    def register_plugin(self, plugin_name: str, config: Optional[Dict[str, Any]] = None) -> bool:
        """
        注册插件（增强依赖管理版本）
        
        Args:
            plugin_name: 插件名称
            config: 插件配置
            
        Returns:
            是否注册成功
        """
        try:
            if config is None:
                config = self.load_plugin_from_directory(plugin_name)
            
            server_name = config.get("server_name", plugin_name)
            
            # 解析并处理依赖关系
            try:
                dependencies = self.resolve_dependencies(plugin_name, config)
                # 先加载所有依赖
                for dep in dependencies:
                    if dep not in self.loaded_plugins:
                        self.logger.info(f"正在加载依赖插件: {dep}")
                        self.register_plugin(dep)
            except MCPPluginDependencyError as e:
                self.logger.error(f"插件 {plugin_name} 依赖解析失败: {e}")
                self._set_plugin_state(plugin_name, self.PLUGIN_STATE_ERROR)
                return False
            
            # 添加到MCP服务器列表
            if server_name not in settings.MCP_SERVERS:
                settings.MCP_SERVERS.append(server_name)
                settings.init_mcp_configs()
            
            # 更新配置
            settings.update_mcp_server_config(server_name, config)
            
            # 保存到已加载插件
            self.loaded_plugins[server_name] = config
            self.plugin_states[server_name] = "loaded"
            
            # 更新依赖关系
            self.plugin_dependencies[plugin_name] = config.get("dependencies", [])
            
            self.logger.info(f"插件 {plugin_name} 注册成功")
            return True
        except Exception as e:
            self.logger.error(f"注册插件 {plugin_name} 失败: {e}")
            return False

    def load_all_plugins(self):
        """加载所有插件"""
        try:
            # 发现所有可用插件
            plugins = self.discover_plugins()
            self.logger.info(f"发现 {len(plugins)} 个插件: {plugins}")
            
            # 加载每个插件
            for plugin_name in plugins:
                try:
                    self.register_plugin(plugin_name)
                    self.logger.info(f"成功加载插件: {plugin_name}")
                except Exception as e:
                    self.logger.error(f"加载插件 {plugin_name} 失败: {e}")
                    
        except Exception as e:
            self.logger.error(f"加载所有插件时发生错误: {e}")
            raise MCPPluginLoadError(f"加载所有插件失败: {str(e)}") from e

    async def browse_marketplace(self) -> List[Dict[str, Any]]:
        """
        浏览插件市场
        
        Returns:
            可用插件列表
        """
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(f"{self.plugin_registry_url}/plugins.json") as response:
                    if response.status == 200:
                        marketplace_data = await response.json()
                        return marketplace_data.get("plugins", [])
                    else:
                        self.logger.error(f"获取插件市场数据失败，状态码: {response.status}")
                        return []
        except Exception as e:
            self.logger.error(f"浏览插件市场失败: {e}")
            return []

    async def install_plugin_from_marketplace(self, plugin_name: str, version: str = "latest") -> bool:
        """
        从插件市场安装插件
        
        Args:
            plugin_name: 插件名称
            version: 插件版本
            
        Returns:
            是否安装成功
        """
        try:
            # 从插件市场获取插件信息
            marketplace_plugins = await self.browse_marketplace()
            plugin_info = next((p for p in marketplace_plugins if p["name"] == plugin_name), None)
            
            if not plugin_info:
                self.logger.error(f"插件市场中未找到插件: {plugin_name}")
                return False
            
            # 检查版本
            if version != "latest":
                available_versions = plugin_info.get("versions", [])
                if version not in available_versions:
                    self.logger.error(f"插件 {plugin_name} 不支持版本: {version}")
                    return False
            
            # 下载插件
            download_url = plugin_info.get("download_url")
            if not download_url:
                self.logger.error(f"插件 {plugin_name} 没有下载地址")
                return False
            
            # 下载并解压插件
            plugin_path = os.path.join(self.plugin_dir, plugin_name)
            if os.path.exists(plugin_path):
                self.logger.warning(f"插件目录 {plugin_path} 已存在，将被覆盖")
                shutil.rmtree(plugin_path)
            
            # 创建插件目录
            os.makedirs(plugin_path, exist_ok=True)
            
            # 下载插件包
            async with aiohttp.ClientSession() as session:
                async with session.get(download_url) as response:
                    if response.status == 200:
                        plugin_zip_path = os.path.join(plugin_path, f"{plugin_name}.zip")
                        with open(plugin_zip_path, 'wb') as f:
                            f.write(await response.read())
                        
                        # 解压插件
                        with zipfile.ZipFile(plugin_zip_path, 'r') as zip_ref:
                            zip_ref.extractall(plugin_path)
                        
                        # 删除zip文件
                        os.remove(plugin_zip_path)
                    else:
                        self.logger.error(f"下载插件 {plugin_name} 失败，状态码: {response.status}")
                        return False
            
            # 注册插件
            success = self.register_plugin(plugin_name)
            
            if success:
                self.logger.info(f"从插件市场成功安装插件: {plugin_name}")
            else:
                self.logger.error(f"安装插件 {plugin_name} 失败")
                
            return success
        except Exception as e:
            self.logger.error(f"从插件市场安装插件 {plugin_name} 失败: {e}")
            return False

    def search_plugins(self, keyword: str) -> List[Dict[str, Any]]:
        """
        搜索插件
        
        Args:
            keyword: 搜索关键词
            
        Returns:
            匹配的插件列表
        """
        try:
            # 搜索已安装的插件
            installed_plugins = []
            for plugin_name, plugin_config in self.loaded_plugins.items():
                if (keyword.lower() in plugin_name.lower() or 
                    keyword.lower() in plugin_config.get("description", "").lower()):
                    plugin_info = plugin_config.copy()
                    plugin_info["name"] = plugin_name
                    plugin_info["installed"] = True
                    installed_plugins.append(plugin_info)
            
            return installed_plugins
        except Exception as e:
            self.logger.error(f"搜索插件失败: {e}")
            return []

    # 添加缺失的方法占位符
    def _get_cached_plugin(self, plugin_name: str) -> Optional[Any]:
        """获取缓存的插件"""
        return self._plugin_cache.get(plugin_name)

    def _cache_plugin(self, plugin_name: str, plugin_data: Any) -> None:
        """缓存插件"""
        self._plugin_cache[plugin_name] = plugin_data

    def _verify_plugin_signature(self, plugin_name: str, content: bytes, signature: str) -> bool:
        """验证插件签名"""
        if not self.plugin_secret_key:
            self.logger.warning(f"插件 {plugin_name} 签名验证跳过：未配置密钥")
            return True
            
        try:
            expected_signature = hmac.new(
                self.plugin_secret_key.encode(),
                content,
                hashlib.sha256
            ).hexdigest()
            return hmac.compare_digest(expected_signature, signature)
        except Exception as e:
            self.logger.error(f"插件 {plugin_name} 签名验证失败: {e}")
            return False

    def _check_plugin_permissions(self, plugin_name: str, config: Dict[str, Any]) -> bool:
        """检查插件权限"""
        # 检查插件是否在允许列表中
        allowed_plugins = os.getenv("ALLOWED_MCP_PLUGINS", "")
        if allowed_plugins:
            allowed_list = [p.strip() for p in allowed_plugins.split(",")]
            if plugin_name not in allowed_list:
                self.logger.warning(f"插件 {plugin_name} 不在允许列表中")
                return False
        
        # 检查插件配置中的权限要求
        required_permissions = config.get("permissions", [])
        if required_permissions:
            self.logger.info(f"插件 {plugin_name} 请求权限: {required_permissions}")
            # 在实际实现中，这里应该检查用户是否授予了这些权限
            
        return True

    def _log_error_with_context(self, error: Exception, context: Dict[str, Any]) -> None:
        """记录带上下文的错误"""
        error_msg = f"Error: {str(error)}"
        if context:
            context_msg = ", ".join([f"{k}: {v}" for k, v in context.items()])
            error_msg += f", Context: {context_msg}"
        self.logger.error(error_msg)

    def _set_plugin_state(self, plugin_name: str, state: str) -> None:
        """
        设置插件状态
        
        Args:
            plugin_name: 插件名称
            state: 状态值
        """
        self.plugin_states[plugin_name] = state
        self.logger.debug(f"插件 {plugin_name} 状态设置为: {state}")

    def unregister_plugin(self, plugin_name: str) -> bool:
        """
        注销插件
        
        Args:
            plugin_name: 插件名称
            
        Returns:
            是否注销成功
        """
        try:
            if plugin_name in self.loaded_plugins:
                # 从已加载插件中移除
                del self.loaded_plugins[plugin_name]
                
                # 从插件状态中移除
                if plugin_name in self.plugin_states:
                    del self.plugin_states[plugin_name]
                
                # 从插件依赖关系中移除
                if plugin_name in self.plugin_dependencies:
                    del self.plugin_dependencies[plugin_name]
                
                # 从设置中移除
                if plugin_name in settings.MCP_SERVERS:
                    settings.MCP_SERVERS.remove(plugin_name)
                
                self.logger.info(f"插件 {plugin_name} 注销成功")
                return True
            else:
                self.logger.warning(f"插件 {plugin_name} 未注册，无法注销")
                return False
        except Exception as e:
            self.logger.error(f"注销插件 {plugin_name} 失败: {e}")
            return False

# 创建全局插件加载器实例
plugin_loader = MCPPluginLoader()